package leetcode.editor.week.week312;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

// 2421. 好路径的数目
// https://leetcode.cn/problems/number-of-good-paths/
public class Solution6191 {
    // https://leetcode.cn/problems/number-of-good-paths/

    // 并查集
    int[] parent;   // 存储父节点

    // 找到父节点的函数
    public int find(int x) {
        if (x != parent[x]) {
            parent[x] = find(parent[x]);
        }
        return parent[x];
    }

    public int numberOfGoodPaths(int[] vals, int[][] edges) {
        int n = vals.length;
        // 表示节点，对节点的val值进行排序，从小到大，先选择小的加入
        Integer[] id = new Integer[n];

        // 初始化
        parent = new int[n];

        // 建图
        List<Integer>[] graph = new ArrayList[n];

        for (int i = 0; i < n; i++) {
            parent[i] = i;
            id[i] = i;
            graph[i] = new ArrayList<>();
        }

        for (int[] edge : edges) {
            int from = edge[0], to = edge[1];
            graph[from].add(to);
            graph[to].add(from);
        }

        // 进行排序
        Arrays.sort(id, (o1, o2) -> vals[o1] - vals[o2]);

        // size[x] 表示节点值等于 vals[x] 的节点个数，如果按照节点值从小到大合并，size[x] 也是连通块内的等于最大节点值的节点个数
        int[] size = new int[n];
        Arrays.fill(size, 1);
        int ans = n;

        for (Integer cur : id) {
            int curVal = vals[cur], curPar = find(cur);
            for (Integer next : graph[cur]) {
                int nextPar = find(next);
                if (curPar == nextPar || curVal < vals[nextPar]) continue;
                if (curVal == vals[nextPar]) {
                    ans += size[curPar] * size[nextPar];
                    size[curPar] += size[nextPar];
                }
                parent[nextPar] = curPar;
            }
        }   

        return ans;
    }
}
